"""**********************************
   Author: yuhengShi
   Time: 2022/1/24 11:23
   Project: pytorch_classification
**********************************"""
from typing import Tuple, Dict

import pandas as pd
from pandas import DataFrame
from tqdm import tqdm

from config import conf
from utils.func import Preprocess


class Vocab:
    def __init__(self, limit=100000):
        """

        :param limit: 数据条数限制
        """
        self.df: DataFrame = pd.read_csv(conf.data_path)[:limit]
        self.wordcount_map = {}
        self.tokens2idx = {}
        self.idx2tokens = {}
        self.preprocessor = Preprocess()

    def create_vocab(self) -> Tuple[Dict, Dict]:
        """

        :return: (idx跟字的map, 字跟idx的map)
        """
        for idx, series in tqdm(self.df.iterrows(), total=len(self.df), desc=f"tokenize level [{conf.tokenize_level}], start"):
            label, text = series
            text: str
            text = text.strip()
            if not text:
                continue
            words = self.preprocessor.tokenize(text)
            self._word_count(words)

        # 根据词频排序
        vocabs = sorted([(k, v) for k, v in self.wordcount_map.items() if k and v >= conf.min_freq],
                        key=lambda x: x[1],
                        reverse=True)[: conf.max_vocab_len]

        # lookup dict
        for index, vocab_data in enumerate(vocabs):
            word, count = vocab_data
            self.tokens2idx[word] = index
            self.idx2tokens[index] = word

        vocabs_len = len(vocabs)
        # 加入UNK以及PAD标记
        self.tokens2idx[conf.UNK] = vocabs_len
        self.tokens2idx[conf.PAD] = vocabs_len + 1

        self.idx2tokens[vocabs_len] = conf.UNK
        self.idx2tokens[vocabs_len + 1] = conf.PAD
        return self.idx2tokens, self.tokens2idx

    def _word_count(self, words):
        for word in words:
            word = word.strip()
            count = self.wordcount_map.get(word)
            num = count + 1 if count else 1
            self.wordcount_map[word] = num
        return self.wordcount_map
